import torch
import torch.nn as nn
import torchvision.models as models

class FCN8s(nn.Module):
    def __init__(self, n_class=21):
        super(FCN8s, self).__init__()
        # 加载预训练的 VGG16 模型
        vgg = models.vgg16(pretrained=True)
        features = list(vgg.features.children())
        self.features1 = nn.Sequential(*features[:17])  # 第一个池化层
        self.features2 = nn.Sequential(*features[17:24])  # 第二个池化层
        self.features3 = nn.Sequential(*features[24:31])  # 第三个池化层
        self.features4 = nn.Sequential(*features[31:])  # 第四个和第五个池化层

        # 全卷积层
        self.fc6 = nn.Conv2d(512, 4096, 7)
        self.relu6 = nn.ReLU(inplace=True)
        self.drop6 = nn.Dropout2d()
        self.fc7 = nn.Conv2d(4096, 4096, 1)
        self.relu7 = nn.ReLU(inplace=True)
        self.drop7 = nn.Dropout2d()

        # 反卷积层
        self.score_fr = nn.Conv2d(4096, n_class, 1)
        self.score_pool3 = nn.Conv2d(512, n_class, 1)  # 修改这里
        self.score_pool4 = nn.Conv2d(512, n_class, 1)

        self.upscore2 = nn.ConvTranspose2d(n_class, n_class, 4, stride=2, bias=False)
        self.upscore8 = nn.ConvTranspose2d(n_class, n_class, 16, stride=8, bias=False)
        self.upscore_pool4 = nn.ConvTranspose2d(n_class, n_class, 2, stride=2, bias=False)

        # 初始化权重
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.ConvTranspose2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        # 特征提取
        pool1 = self.features1(x)
        pool2 = self.features2(pool1)
        pool3 = self.features3(pool2)
        pool4 = self.features4(pool3)

        fc6 = self.relu6(self.fc6(pool4))
        fc6 = self.drop6(fc6)
        fc7 = self.relu7(self.fc7(fc6))
        fc7 = self.drop7(fc7)

        score_fr = self.score_fr(fc7)
        score_pool4 = self.score_pool4(pool4)
        score_pool3 = self.score_pool3(pool3)

        upscore2 = self.upscore2(score_fr)

        upscore2,score_pool4 = self._crop(upscore2,score_pool4)
        # print("upscore2",upscore2.size())
        # print("score_pool4", score_pool4.size())
        fuse_pool4 = upscore2 + score_pool4
        upscore_pool4 = self.upscore_pool4(fuse_pool4)
        upscore_pool4, score_pool3 = self._crop(upscore_pool4, score_pool3)
        fuse_pool3 = upscore_pool4 + score_pool3
        # print("fuse_pool3", fuse_pool3.size())
        outputs = self.upscore8(fuse_pool3)
        # print("outputs", outputs.size())

        return outputs

    # def _crop(self, tensor1, tensor2):
    #     _, _, h1, w1 = tensor1.size()
    #     _, _, h2, w2 = tensor2.size()
    #     return tensor1[:, :, :h2, :w2]

    def _crop(self, tensor1, tensor2):
        # 获取两个张量的高度和宽度
        _, _, h1, w1 = tensor1.size()
        _, _, h2, w2 = tensor2.size()

        # 确定较小的高度和宽度
        target_height = min(h1, h2)
        target_width = min(w1, w2)

        # 计算裁剪的起始位置
        start_h1 = (h1 - target_height) // 2
        start_w1 = (w1 - target_width) // 2
        start_h2 = (h2 - target_height) // 2
        start_w2 = (w2 - target_width) // 2

        # 裁剪张量
        cropped_tensor1 = tensor1[:, :, start_h1:start_h1 + target_height, start_w1:start_w1 + target_width]
        cropped_tensor2 = tensor2[:, :, start_h2:start_h2 + target_height, start_w2:start_w2 + target_width]

        return cropped_tensor1, cropped_tensor2